from pathlib import Path
import subprocess
import shutil
import sys
import logging
import os

import features.angr_adapter as adapter

def run_angr_disasm(binary_path, asm_path, block_path):
    """
    运行 angr_disasm.py 脚本，生成二进制的asm文件。
    
    参数:
    - binary_path (str): 二进制文件路径。
    - asm_path (str): 反汇编输出文件路径。
    - block_path (str): 基本块输出路径。
    """
    script_dir = Path(__file__).parent
    angr_disasm_script = script_dir / 'angr_disasm.py'
    subprocess.run([
        sys.executable, str(angr_disasm_script),
        '--binary_path', binary_path,
        '--asm_path', asm_path,
        '--block_path', block_path,
    ], check=True)

def run_decompile(src_folder, out_folder, ida_path, script_path):
    """
    运行 run.py 脚本，调用inline_determine.py,生成反汇编文件，基本块信息文件，反编译文件和内联检测文件。
    
    参数:
    - src_folder (str): 源文件夹路径。
    - out_folder (str): 输出文件夹路径。
    - ida_path (str): IDA Pro 工具路径。
    - script_path (str): IDA 脚本路径。
    """
    script_dir = Path(__file__).parent
    run_script = script_dir / 'run.py'
    subprocess.run([
        sys.executable, str(run_script),
        '--src_folder', src_folder,
        '--out_folder', out_folder,
        '--ida_path', ida_path,
        '--script_path', script_path,
        '--run_type', 'decompile'
    ], check=True)

def run_asm_parser(asm_file, code_file, inline_file, block_file, output_file, missing_file):
    """
    运行 asm_parser.py 脚本，生成asm注解。
    
    参数:
    - asm_file (str): ASM 文件路径。
    - code_file (str): 反汇编代码文件路径。
    - inline_file (str): 内联反汇编文件路径。
    - block_file (str): 块信息文件路径。
    - output_file (str): 结果 JSON 文件路径。
    - missing_file (str): 地址不匹配 JSON 文件路径。
    """
    script_dir = Path(__file__).parent
    asm_parser_script = script_dir / 'asm_parser.py'
    subprocess.run([
        sys.executable, str(asm_parser_script),
        '--asm_file', asm_file,
        '--code_file', code_file,
        '--inline_file', inline_file,
        '--block_file', block_file,
        '--output_file', output_file,
        '--missing_file', missing_file
    ], check=True)
    
# def run_angr_adapter(binary_file_path, binary_output_dir, graph_types=['CG']):
#     """
#     运行 angr_adapter.batch_processing 生成控制流图（CG）。
    
#     参数:
#     - binary_file_path (str): 单个二进制文件路径。
#     - binary_output_dir (str): 对应的输出目录，用于存储 CG 文件。
#     - graph_types (list): 生成的图类型，例如 ['CFG', 'VFG']，默认为 ['CG']。
#     """

#     # 调用 angr_adapter.batch_processing 生成 CG
#     adapter.batch_processing(
#         binary_file_path,
#         binary_output_dir,
#         graph_types
#     )
def run_ida_cg(src_folder, out_folder, ida_path, script_path):
    """
    运行 run.py 脚本，调用ida_cg.py,生成生成控制流图（CG）。
    
    参数:
    - src_folder (str): 源文件夹路径。
    - out_folder (str): 输出文件夹路径。
    - ida_path (str): IDA Pro 工具路径。
    - script_path (str): IDA 脚本路径。
    """
    script_dir = Path(__file__).parent
    run_script = script_dir / 'run.py'
    subprocess.run([
        sys.executable, str(run_script),
        '--src_folder', src_folder,
        '--out_folder', out_folder,
        '--ida_path', ida_path,
        '--script_path', script_path,
    ], check=True)
    
def run_decomp_parser(decompiled_code_path, inline_disassembly_path, cg_graph_path, output_json_path):
    """
    运行 decomp_parser.py 脚本，解析反编译结果。
    
    参数:
    - decompiled_code_path (str): 反编译文件路径。
    - inline_disassembly_path (str): 内联反汇编文件路径。
    - cg_graph_path (str): 生成的图json文件路径。
    - output_json_path (str): 结果 JSON 文件路径。
    """
    script_dir = Path(__file__).parent
    decomp_parser_script = script_dir / 'decomp_parser.py'
    subprocess.run([
        sys.executable, str(decomp_parser_script),
        '--decompiled_code_path', decompiled_code_path,
        '--inline_disassembly_path', inline_disassembly_path,
        '--cg_graph_path', cg_graph_path,
        '--output_json_path', output_json_path
    ], check=True)

def run_all(input_dir, output_dir, ida_path, script_path):
    """
    对指定目录下的所有二进制文件依次运行 angr_disasm.py、run.py 和 asm_parser.py，
    并将生成的结果存入另一个指定位置的以目录下文件名命名的文件夹中。
    
    参数:
    - input_dir (str): 输入目录，包含待处理的二进制文件。
    - output_dir (str): 输出目录，用于存储处理结果。
    - ida_path (str): IDA Pro 工具的路径。
    - script_path (str): IDA 脚本（inline_determine.py）的路径。
    """
    # 配置日志
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

    input_dir = Path(input_dir)
    output_dir = Path(output_dir)
    ida_path = Path(ida_path)
    script_path = Path(script_path)

    # 验证路径
    if not input_dir.is_dir():
        logging.error(f"输入目录 {input_dir} 不存在或不是一个目录。")
        sys.exit(1)
    if not ida_path.exists():
        logging.error(f"IDA Pro 工具路径 {ida_path} 不存在。")
        sys.exit(1)
    if not script_path.exists():
        logging.error(f"IDA 脚本路径 {script_path} 不存在。")
        sys.exit(1)

    output_dir.mkdir(parents=True, exist_ok=True)

    # 遍历输入目录中的所有文件
    for binary_file in input_dir.iterdir():
        if binary_file.is_file():
            binary_name = binary_file.stem
            logging.info(f"处理文件: {binary_file.name}")

            # 创建每个二进制文件的输出子目录
            binary_output_dir = output_dir / binary_name
            binary_output_dir.mkdir(parents=True, exist_ok=True)

            try:
                # Step 1: 运行 angr_disasm.py
                disasm_output_path = binary_output_dir / f"{binary_name}.asm"
                block_output_path = binary_output_dir / f"{binary_name}_block_info.txt"
                logging.info(f"运行 angr_disasm.py，asm输出路径: {disasm_output_path}, block输出路径:{block_output_path}")
                run_angr_disasm(str(binary_file.resolve()), str(disasm_output_path), str(block_output_path))

                # Step 2: 运行 run.py
                temp_src_folder = binary_output_dir / "temp_src"
                temp_src_folder.mkdir(parents=True, exist_ok=True)
                temp_copy = temp_src_folder / binary_file.name
                os.chmod(binary_file, 0o777)
                shutil.copy2(binary_file, temp_copy)

                logging.info(f"运行 run.py，src_folder: {temp_src_folder}, out_folder: {binary_output_dir}")
                run_decompile(
                    str(temp_src_folder.resolve()),
                    str(binary_output_dir.resolve()),
                    str(ida_path.resolve()),
                    str(script_path.resolve())
                )

                # Step 3: 运行 asm_parser.py
                asm_file = binary_output_dir / f"{binary_name}.asm"
                code_file = binary_output_dir / f"{binary_name}_assembly_code.txt"
                inline_file = binary_output_dir / f"{binary_name}_inline_disassembly.txt"
                block_file = binary_output_dir / f"{binary_name}_block_info.txt"
                result_file = binary_output_dir / "result.json"
                missing_file = binary_output_dir / "address_mismatch.json"

                required_files = [asm_file, code_file, inline_file, block_file]
                missing_files = [f for f in required_files if not f.exists()]
                if missing_files:
                    logging.warning(f"缺少必要的文件: {missing_files}，跳过 asm_parser.py 的运行。")
                else:
                    logging.info(f"运行 asm_parser.py，asm_file: {asm_file}, code_file: {code_file}, "
                                 f"inline_file: {inline_file}, block_file: {block_file}")
                    run_asm_parser(
                        str(asm_file.resolve()),
                        str(code_file.resolve()),
                        str(inline_file.resolve()),
                        str(block_file.resolve()),
                        str(result_file.resolve()),
                        str(missing_file.resolve())
                    )
                    
                # Step 4: 运行 angr_adapter.batch_processing 生成控制流图（CG）
                # logging.info("运行 angr_adapter.batch_processing 生成控制流图（CG）")
                # run_angr_adapter(
                #     binary_file_path=str(temp_src_folder.resolve()),
                #     binary_output_dir=str(output_dir.resolve()),
                #     graph_types='cg'
                # )
                # logging.info("angr_adapter.batch_processing 完成\n")
                
                logging.info("运行 ida_cg.py 生成控制流图（CG）")
                cg_script_path = str(script_path).replace('inline_determine.py', 'ida_cg.py')
                run_decompile(
                    str(temp_src_folder.resolve()),
                    str(binary_output_dir.resolve()),
                    str(ida_path.resolve()),
                    str(cg_script_path)
                )
                logging.info("IDA CG生成 完成\n")
                
                # Step 5: 运行 decomp_parser.py
                cg_file = binary_output_dir / "cg_graph.json"
                decompiled_file = binary_output_dir / f"{binary_name}_decompiled_code.c"
                
                derequired_files = [cg_file, decompiled_file, inline_file]
                demissing_files = [f for f in derequired_files if not f.exists()]
                if missing_files:
                    logging.warning(f"缺少必要的文件: {demissing_files}，跳过 asm_parser.py 的运行。")
                else:
                    output_json_path = binary_output_dir / "decompile_result.json"
                    logging.info(f"运行 decomp_parser.py，decompiled_code_path: {decompiled_file}, "
                                 f"inline_disassembly_path: {inline_file}, "
                                 f"cg_graph_path: {cg_file}, output_json_path: {output_json_path}")
                    run_decomp_parser(
                        decompiled_code_path=str(decompiled_file.resolve()),
                        inline_disassembly_path=str(inline_file.resolve()),
                        cg_graph_path=str(cg_file.resolve()),
                        output_json_path=str(output_json_path.resolve())
                    )

                logging.info(f"完成处理: {binary_file.name}\n")

            except subprocess.CalledProcessError as e:
                logging.error(f"处理文件 {binary_file.name} 时出错: {e}")
            finally:
                # 清理临时源文件夹
                if temp_src_folder.exists():
                    shutil.rmtree(temp_src_folder)

                    
if __name__ == "__main__":
    input_dir = r'D:\Users\Administrator\Desktop\study\bad'
    output_dir = r'D:\Users\Administrator\Desktop\study\项目\IDA\out'
    ida_path = r'D:\Users\Administrator\Desktop\study\IDA_Pro_7.5_SP3\IDA_Pro_7.5_SP3'
    script_path = r'D:\Users\Administrator\Desktop\study\项目\IDA\inline_determine.py'
    run_all(input_dir, output_dir, ida_path, script_path)
